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

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

using namespace std;
using namespace htqc;

//
// type definitions
//

enum TrimSide
{
    TRIM_FROM_HEAD = 1,
    TRIM_FROM_TAIL = 1 << 1
};

string trim_side_to_string(TrimSide side)
{
    if (side == TRIM_FROM_HEAD)
        return "head";
    else if (side == TRIM_FROM_TAIL)
        return "tail";
    else if (side == (TRIM_FROM_HEAD | TRIM_FROM_TAIL))
        return "both";
    else
    {
        cerr << "ERROR: invalid side enum value: " << int(side) << endl;
        exit(EXIT_FAILURE);
    }
}

//
// command line options
//
#define DEFAULT_KMER_SIZE 5
#define DEFAULT_QUAL_CUT 20

#define OPTION_GROUP_TRIM "Trimming Options"

unsigned char OPT_trim_side = TRIM_FROM_HEAD | TRIM_FROM_TAIL;
string OPT_trim_side_STR("both");
htio2::Option OPT_trim_side_ENTRY("side", 'S', OPTION_GROUP_TRIM,
                                  &OPT_trim_side_STR, htio2::Option::FLAG_NONE,
                                  "Trim is performed on which side.", "head|tail|both");

size_t OPT_word_size = DEFAULT_KMER_SIZE;
htio2::Option OPT_word_size_ENTRY("word-size", 'W', OPTION_GROUP_TRIM,
                                  &OPT_word_size, htio2::Option::FLAG_NONE,
                                  "Number of continuous high quality bases.", "INT");

int OPT_qual_cut = DEFAULT_QUAL_CUT;
htio2::Option OPT_qual_cut_ENTRY("quality", 'Q', OPTION_GROUP_TRIM,
                                 &OPT_qual_cut, htio2::Option::FLAG_NONE,
                                 "Quality cutoff.", "INT");

//
// subs
//

void parse_options(int argc, char** argv)
{
    // we are always processing single-end
    OPT_se = 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_trim_side_ENTRY);
    opt_psr.add_option(OPT_qual_cut_ENTRY);
    opt_psr.add_option(OPT_word_size_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] << " - remove bad bases from reads head or tail"
            << 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();

    if (OPT_word_size < 1 || OPT_word_size > 16)
    {
        cerr << "ERROR: invalid kmer size: " << OPT_word_size << ", must between 1 - 16" << endl;
        exit(EXIT_FAILURE);
    }

    if (OPT_trim_side_STR == "head")
        OPT_trim_side = TRIM_FROM_HEAD;
    else if (OPT_trim_side_STR == "tail")
        OPT_trim_side = TRIM_FROM_TAIL;
    else if (OPT_trim_side_STR == "both")
        OPT_trim_side = TRIM_FROM_HEAD | TRIM_FROM_TAIL;
    else
    {
        cerr << "ERROR: invalid trim side: " << OPT_trim_side_STR << ", must be head|tail|both" << endl;
        exit(EXIT_FAILURE);
    }

    ensure_encode();
}

void show_options()
{
    show_files_in();
    show_prefix_out();
    cout << "# trim side: " << OPT_trim_side_STR << endl;
    cout << "# word size: " << OPT_word_size << endl;
    cout << "# quality cutoff: " << OPT_qual_cut << endl;
    cout << "# encode: " << to_string(OPT_encode) << endl;
}

string to_str(uint16_t kmer)
{
    string re;
    for (size_t i = 0; i < OPT_word_size; i++)
    {
        if (kmer & (1 << i))
            re.push_back('1');
        else
            re.push_back('0');
    }
    return re;
}

void execute()
{
    // 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));


    // generate kmer that is perfect
    uint16_t perfect = 0;
    for (size_t i = 0; i < OPT_word_size; i++)
        perfect |= 1 << i;

    //
    // traverse input
    //
    size_t n = 0;
    htio2::FastqSeq curr_seq;
    htio2::FastqSeq curr_seq_sub;

    htqc::MultiSeqFileSE fh_in(OPT_files_in, OPT_compress);

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

        ssize_t len = curr_seq.length();
        ssize_t i_start = 0;
        ssize_t i_end = len - 1;

        // sequence must be long enough
        if (len < OPT_word_size * 2)
        {
            i_end = i_start - 1;
            goto TRIM_DONE;
        }

        // trim from head
        if (OPT_trim_side & TRIM_FROM_HEAD)
        {
            // the first kmer
            uint16_t kmer = 0;
            for (size_t i = 0; i < OPT_word_size; i++)
            {
                int32_t qual = htio2::decode_quality(curr_seq.quality[i], OPT_encode);
                bool curr_ok = (OPT_qual_cut <= qual);
                kmer |= curr_ok << i;
            }

            // traverse all remaining kmers
            if (kmer != perfect)
            {
                while (true)
                {
                    if (i_start == len - OPT_word_size) break;
                    i_start++;

                    int32_t qual = htio2::decode_quality(curr_seq.quality[i_start + OPT_word_size - 1], OPT_encode);
                    bool curr_ok = (OPT_qual_cut <= qual);
                    kmer >>= 1;
                    kmer |= curr_ok << (OPT_word_size - 1);
                    kmer &= perfect;
                    if (kmer == perfect) break;
                }
            }
        }

        // trim from tail
        if (OPT_trim_side & TRIM_FROM_TAIL)
        {
            // the last kmer
            uint16_t kmer = 0;
            for (size_t i = 0; i < OPT_word_size; i++)
            {
                int32_t qual = htio2::decode_quality(curr_seq.quality[len - OPT_word_size + i], OPT_encode);
                bool curr_ok = (OPT_qual_cut <= qual);
                kmer |= curr_ok << i;
            }

            // traverse all remaining kmers
            if (kmer != perfect)
            {
                while (true)
                {
                    if (i_end == i_start - 1) break;
                    i_end--;
                    int qual = htio2::decode_quality(curr_seq.quality[i_end - OPT_word_size + 1], OPT_encode);
                    bool curr_ok = (OPT_qual_cut <= qual);
                    kmer <<= 1;
                    kmer |= curr_ok;
                    kmer &= perfect;
                    if (kmer == perfect) break;
                }
            }
        }

        // write output
TRIM_DONE:
        int trim_len = i_end - i_start + 1;
        curr_seq.subseq(curr_seq_sub, i_start, trim_len);
        fh_out->write_seq(curr_seq_sub);

        if (!OPT_quiet && (n % LOG_BLOCK == 0))
            cout << "  " << n << " reads" << endl;
    }


    if (!OPT_quiet)
        cout << "  " << n << " reads" << endl;
}
