#include "htio2/FastqIO.h"
#include "htio2/FastqSeq.h"
#include "htio2/PairedEndIO.h"
#include "htio2/Kmer.h"
#include "htio2/FastaIO.h"
#include "htio2/QualityUtil.h"

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

#include "htio2/JUCE-3.0.8/JuceHeader.h"

// we need paired-end reads in fwd-rev style
#define PE_STRAND_FILE OPT_pe_strand1(), OPT_pe_strand2()
#define PE_STRAND_OBJ htio2::STRAND_FWD, htio2::STRAND_REV

using namespace std;
using namespace htqc;
using namespace htio2::juce;

string OPT_dir_out;
htio2::Option OPT_dir_out_ENTRY("out", 'o', OPTION_GROUP_IO,
                                &OPT_dir_out, htio2::Option::FLAG_NONE,
                                "Output directory.", "DIR");

string OPT_file_primer;
htio2::Option OPT_file_primer_ENTRY("primer", '\0', OPTION_GROUP_IO,
                                    &OPT_file_primer, htio2::Option::FLAG_NONE,
                                    "Primers in FASTA format. IDs of left primers should begin with \"left\", and IDs for right primers should begin with \"right\".",
                                    "FILE");

#define OPTION_GROUP_ALGO "Primer Trimming Parameters"
bool OPT_no_revcom;
htio2::Option OPT_no_revcom_ENTRY("no-revcom", '\0', OPTION_GROUP_ALGO,
                                  &OPT_no_revcom, htio2::Option::FLAG_NONE,
                                  "Do not check primers in reverse complement strand.");

bool OPT_no_trim;
htio2::Option OPT_no_trim_ENTRY("no-trim", '\0', OPTION_GROUP_ALGO,
                                &OPT_no_trim, htio2::Option::FLAG_NONE,
                                "Pick primer-containing reads and do not trim primer.");

int OPT_mismatch_high = 3;
htio2::Option OPT_mismatch_high_ENTRY("mismatch-high", '\0', OPTION_GROUP_ALGO,
                                      &OPT_mismatch_high, htio2::Option::FLAG_NONE,
                                      "Allowed number of mismatches on high-quality bases.",
                                      "INT");

int OPT_mismatch = 6;
htio2::Option OPT_mismatch_ENTRY("mismatch", '\0', OPTION_GROUP_ALGO,
                                 &OPT_mismatch, htio2::Option::FLAG_NONE,
                                 "Allowed number of mismatches.",
                                 "INT");

int OPT_tail_miss = 5;
htio2::Option OPT_tail_miss_ENTRY("tail-miss", '\0', OPTION_GROUP_ALGO,
                                  &OPT_tail_miss, htio2::Option::FLAG_NONE,
                                  "Allowed length of missing bases on primer 5'-end.",
                                  "INT");

int OPT_cut_qual = 10;
htio2::Option OPT_cut_qual_ENTRY("cut-qual", '\0', OPTION_GROUP_ALGO,
                                 &OPT_cut_qual, htio2::Option::FLAG_NONE,
                                 "Quality cutoff.",
                                 "INT");



// primer data
typedef map<string, PrimerGroup> AllPrimerType;
AllPrimerType primers;

void parse_options(int argc, char** argv)
{
    htio2::OptionParser opt_psr;
    opt_psr.add_option(OPT_files_in_ENTRY);
    opt_psr.add_option(OPT_dir_out_ENTRY);
    opt_psr.add_option(OPT_file_primer_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_pe_strand_ENTRY);
    opt_psr.add_option(OPT_encode_ENTRY);
    opt_psr.add_option(OPT_mask_ENTRY);

    opt_psr.add_option(OPT_no_revcom_ENTRY);
    opt_psr.add_option(OPT_no_trim_ENTRY);
    opt_psr.add_option(OPT_cut_qual_ENTRY);
    opt_psr.add_option(OPT_tail_miss_ENTRY);
    opt_psr.add_option(OPT_mismatch_ENTRY);
    opt_psr.add_option(OPT_mismatch_high_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 primer sequence from reads"
             << 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();

    if (!OPT_dir_out.length())
    {
        cerr << "ERROR: output directory not specified!" << endl;
        exit(EXIT_FAILURE);
    }

    if (!OPT_file_primer.length())
    {
        cerr << "ERROR: primer file not specified!" << endl;
        exit(EXIT_FAILURE);
    }
}

void show_options()
{
    show_se_pe();
    show_files_in();
    cout << "# output dir: " << OPT_dir_out << endl;
    cout << "# primers: " << OPT_file_primer << endl;
    cout << "# allowed mismatch for high-quality bases: " << OPT_mismatch_high << endl;
    cout << "# allowed mismatch: " << OPT_mismatch << endl;
    cout << "# cutoff for low-quality bases: " << OPT_cut_qual << endl;
    show_encode();
}

size_t kmer_size = 0;

void read_primers()
{
    // first pass
    // find minimum length
    size_t min_len = 65535;
    {
        htio2::FastaIO IN(OPT_file_primer, htio2::READ);
        htio2::SimpleSeq seq;

        while (IN.next_seq(seq))
        {
            if (seq.seq.length() < min_len)
                min_len = seq.seq.length();
        }
    }

    // calculate kmer size
    kmer_size = (min_len - OPT_mismatch_high) / (OPT_mismatch_high + 1);
    if (kmer_size < 4) kmer_size = 4;
    if (!OPT_quiet) cout << "  shortest primer: " << min_len << ", alignment kmer size: " << kmer_size << endl;

    // read primers
    htio2::FastaIO IN(OPT_file_primer, htio2::READ);
    htio2::SimpleSeq seq;

    if (!OPT_quiet)cout << "read primers" << endl;
    while (IN.next_seq(seq))
    {
        if (!seq.desc.length())
        {
            cerr << "primer \"" << seq.id << "\" has no description" << endl;
            exit(EXIT_FAILURE);
        }

        // find or create the corresponding primer group
        AllPrimerType::iterator primer_map_i = primers.find(seq.desc);
        if (primer_map_i == primers.end())
            primer_map_i = primers.insert(make_pair(seq.desc, PrimerGroup(seq.desc))).first;
        PrimerGroup& primer_group = primer_map_i->second;

        // add primer to primer group
        if (seq.id.find("left") == 0)
            primer_group.primers_left.push_back(Primer(seq.id, seq.seq, kmer_size));
        else if (seq.id.find("right") == 0)
            primer_group.primers_right.push_back(Primer(seq.id, seq.seq, kmer_size));
        else
        {
            cerr << "ERROR: primer file has sequence with invalid ID: " << seq.id << endl
                 << "IDs must begin with \"left\" or \"right\"" << endl;
            exit(EXIT_FAILURE);
        }
    }

    // check primer number
    if (!OPT_quiet)
        cout << "  " << primers.size() << " primer groups" << endl;

    for (AllPrimerType::const_iterator i = primers.begin(); i != primers.end(); ++i)
    {
        const PrimerGroup& group = i->second;

        if (!group.primers_left.size())
        {
            cerr << "ERROR: group " << group.name << " has no left primer" << endl;
            exit (EXIT_FAILURE);
        }

        if (!group.primers_right.size())
        {
            cerr << "ERROR: group " << group.name << " has no right primer" << endl;
            exit (EXIT_FAILURE);
        }

        if (!OPT_quiet)
            cout << "    " << group.name << " left " << group.primers_left.size() << ", right " << group.primers_right.size() << endl;
    }
}



#define LEFT_START_OK  primer_start <= OPT_tail_miss
#define LEFT_END_OK primer_end + 1 == curr_primer_len
#define RIGHT_START_OK  primer_start == 0
#define RIGHT_END_OK curr_primer_len - primer_end - 1 <= OPT_tail_miss

size_t n_seq = 0;
size_t n_accept = 0;

void trim_se()
{
    // create output file handles
    map<string, htio2::SeqIO::Ptr> fh_out_map;
    for (map<string, PrimerGroup>::const_iterator i = primers.begin(); i != primers.end(); i++)
    {
        const string& name = i->first;
        File prefix_out = File(OPT_dir_out).getChildFile(String(name));
        string file_out;
        add_fastq_suffix_se(prefix_out.getFullPathName().toStdString(), 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));
        fh_out_map.insert(make_pair(name, fh_out));
    }

    // create file handle for failed ones
    htio2::SeqIO::Ptr fh_fail(0);
    {
        File prefix_fail = File(OPT_dir_out).getChildFile("PRIMER_FAIL");
        string file_fail;
        add_fastq_suffix_se(prefix_fail.getFullPathName().toStdString(), 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;
    htio2::EncodedSeq curr_enc;
    KmerList curr_kmer;
    vector<int16_t> curr_qual;

    htqc::MultiSeqFileSE fh_in(OPT_files_in, OPT_compress);

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

        // preprocess current read
        htio2::encode_nt5(curr_seq.seq, curr_enc);
        htio2::gen_kmer_nt5(curr_enc, curr_kmer, kmer_size);
        htio2::decode_quality(curr_seq.quality, curr_qual, OPT_encode);

        // traverse primer groups
        bool matched = false;

        for (AllPrimerType::const_iterator i = primers.begin(); i != primers.end(); i++)
        {
            const string& primer_group_name = i->first;
            const PrimerGroup& primer_group = i->second;

            PosT read_start = 0;
            PosT read_end = 0;
            PosT primer_start = 0;
            PosT primer_end = 0;

            bool left_fwd_ok = false;
            bool right_fwd_ok = false;
            bool left_rev_ok = false;
            bool right_rev_ok = false;
            PosT trim_fwd_from = 0;
            PosT trim_fwd_to = 0;
            PosT trim_rev_from = 0;
            PosT trim_rev_to = 0;

            //
            // compare with primers
            //

            size_t num_left = primer_group.primers_left.size();
            size_t num_right = primer_group.primers_right.size();

            // left primer
            for (size_t i = 0; i < num_left; i++)
            {
                const size_t curr_primer_len = primer_group.primers_left[i].length();
                bool aligned = primer_group.primers_left[i].align_fwd(curr_seq.seq, curr_qual, curr_kmer,
                                                                      OPT_cut_qual, OPT_mismatch_high, OPT_mismatch,
                                                                      read_start, read_end, primer_start, primer_end);
                if (aligned && LEFT_START_OK && LEFT_END_OK)
                {
                    left_fwd_ok = true;
                    trim_fwd_from = read_end + 1;
                    break;
                }
            }

            // right primer revcom
            for (size_t i = 0; i < num_right; i++)
            {
                const size_t curr_primer_len = primer_group.primers_right[i].length();
                bool aligned = primer_group.primers_right[i].align_rev(curr_seq.seq, curr_qual, curr_kmer,
                                                                       OPT_cut_qual, OPT_mismatch_high, OPT_mismatch,
                                                                       read_start, read_end, primer_start, primer_end);
                if (aligned && RIGHT_START_OK && RIGHT_END_OK)
                {
                    right_fwd_ok = true;
                    trim_fwd_to = read_start - 1;
                    break;
                }
            }

            // compare with reverse complement strand
            if (!OPT_no_revcom)
            {
                // right primer
                for (size_t i = 0; i < num_right; i++)
                {
                    const size_t curr_primer_len = primer_group.primers_right[i].length();
                    bool aligned = primer_group.primers_right[i].align_fwd(curr_seq.seq, curr_qual, curr_kmer,
                                                                           OPT_cut_qual, OPT_mismatch_high, OPT_mismatch,
                                                                           read_start, read_end, primer_start, primer_end);
                    // aligned on left, using left rules
                    if (aligned && LEFT_START_OK && LEFT_END_OK)
                    {
                        right_rev_ok = true;
                        trim_rev_from = read_end + 1;
                        break;
                    }
                }

                // left primer revcom
                for (size_t i = 0; i < num_left; i++)
                {
                    const size_t curr_primer_len = primer_group.primers_left[i].length();
                    bool aligned = primer_group.primers_left[i].align_rev(curr_seq.seq, curr_qual, curr_kmer,
                                                                          OPT_cut_qual, OPT_mismatch_high, OPT_mismatch,
                                                                          read_start, read_end, primer_start, primer_end);
                    // aligned on right, using right rules
                    if (aligned && RIGHT_START_OK && RIGHT_END_OK)
                    {
                        left_rev_ok = true;
                        trim_rev_to = read_start - 1;
                        break;
                    }
                }
            }

            // write output
            if ((!num_left || left_fwd_ok) && (!num_right || right_fwd_ok))
            {
                n_accept++;
                matched = true;
                if (OPT_no_trim)
                {
                    fh_out_map[primer_group_name]->write_seq(curr_seq);
                }
                else
                {
                    htio2::FastqSeq trimmed_seq;
                    curr_seq.subseq(trimmed_seq, trim_fwd_from, trim_fwd_to - trim_fwd_from + 1);
                    fh_out_map[primer_group_name]->write_seq(trimmed_seq);
                }
            }
            else if ((!num_left || left_rev_ok) && (!num_right || right_rev_ok))
            {
                n_accept++;
                matched = true;
                if (OPT_no_trim)
                {
                    fh_out_map[primer_group_name]->write_seq(curr_seq);
                }
                else
                {
                    htio2::FastqSeq trimmed_seq;
                    curr_seq.subseq(trimmed_seq, trim_rev_from, trim_rev_to - trim_rev_from + 1);
                    fh_out_map[primer_group_name]->write_seq(trimmed_seq);
                }
            }

            if (matched) break;
        } // primer group cycle

        if (!matched)
            fh_fail->write_seq(curr_seq);

        if (!OPT_quiet && n_seq % LOG_BLOCK == 0)
            cout << "  " << n_seq << " reads, " << n_accept << " trimmed" << endl;
    } // sequence cycle


    if (!OPT_quiet)
        cout << "done: " << n_seq << " reads, " << n_accept << " trimmed" << endl;
}

void do_one_pe(const htio2::FastqSeq& seq1,
               const htio2::FastqSeq& seq2,
               map<string, htio2::PairedEndIO::Ptr>& fh_out,
               htio2::PairedEndIO* fh_fail)
{
    n_seq++;

    static htio2::EncodedSeq enc1;
    static htio2::EncodedSeq enc2;
    static KmerList klist1;
    static KmerList klist2;
    static vector<int16_t> qlist1;
    static vector<int16_t> qlist2;

    PosT read_start = 0;
    PosT read_end = 0;
    PosT primer_start = 0;
    PosT primer_end = 0;

    bool left_fwd_ok = false;
    bool right_fwd_ok = false;
    bool left_rev_ok = false;
    bool right_rev_ok = false;
    PosT trim_fwd_r1 = 0;
    PosT trim_fwd_r2 = 0;
    PosT trim_rev_r1 = 0;
    PosT trim_rev_r2 = 0;

    // preprocess current read
    htio2::encode_nt5(seq1.seq, enc1);
    htio2::encode_nt5(seq2.seq, enc2);
    htio2::gen_kmer_nt5(enc1, klist1, kmer_size);
    htio2::gen_kmer_nt5(enc2, klist2, kmer_size);
    htio2::decode_quality(seq1.quality, qlist1, OPT_encode);
    htio2::decode_quality(seq2.quality, qlist2, OPT_encode);

    // traverse primer groups
    bool matched = false;
    for (AllPrimerType::const_iterator i = primers.begin(); i != primers.end(); i++)
    {
        const string& primer_group_name = i->first;
        const PrimerGroup& primer_group = i->second;

        size_t n_primer_left = primer_group.primers_left.size();
        size_t n_primer_right = primer_group.primers_right.size();

        // left on read 1
        for (size_t i = 0; i < n_primer_left; i++)
        {
            const size_t curr_primer_len = primer_group.primers_left[i].length();
            bool aligned = primer_group.primers_left[i].align_fwd(seq1.seq, qlist1, klist1,
                                                                  OPT_cut_qual, OPT_mismatch_high, OPT_mismatch,
                                                                  read_start, read_end, primer_start, primer_end);

            if (aligned && LEFT_START_OK && LEFT_END_OK)
            {
                left_fwd_ok = true;
                trim_fwd_r1 = read_end + 1;
                break;
            }
        }

        // right on read 2
        for (size_t i = 0; i < n_primer_right; i++)
        {
            const size_t curr_primer_len = primer_group.primers_right[i].length();
            bool aligned = primer_group.primers_right[i].align_fwd(seq2.seq, qlist2, klist2,
                                                                   OPT_cut_qual, OPT_mismatch_high, OPT_mismatch,
                                                                   read_start, read_end, primer_start, primer_end);
            if (aligned && LEFT_START_OK && LEFT_END_OK)
            {
                right_fwd_ok = true;
                trim_fwd_r2 = read_end + 1;
                break;
            }
        }

        // revcom
        if (!OPT_no_revcom)
        {
            // align right revcom on read 1
            for (size_t i = 0; i < n_primer_right; i++)
            {
                const size_t curr_primer_len = primer_group.primers_right[i].length();
                bool aligned = primer_group.primers_right[i].align_rev(seq1.seq, qlist1, klist1,
                                                                       OPT_cut_qual, OPT_mismatch_high, OPT_mismatch,
                                                                       read_start, read_end, primer_start, primer_end);
                if (aligned && LEFT_START_OK && LEFT_END_OK)
                {
                    right_rev_ok = true;
                    trim_rev_r1 = read_end + 1;
                    break;
                }
            }

            // align left revcom on read 2
            for (size_t i = 0; i < n_primer_left; i++)
            {
                const size_t curr_primer_len = primer_group.primers_left[i].length();
                bool aligned = primer_group.primers_left[i].align_rev(seq2.seq, qlist2, klist2,
                                                                      OPT_cut_qual, OPT_mismatch_high, OPT_mismatch,
                                                                      read_start, read_end, primer_start, primer_end);
                if (aligned && LEFT_START_OK && LEFT_END_OK)
                {
                    left_rev_ok = true;
                    trim_rev_r2 = read_end + 1;
                    break;
                }
            }

        }

        // fwd ok
        if ((!n_primer_left || left_fwd_ok)
                &&
                (!n_primer_right || right_fwd_ok))
        {
            n_accept++;
            matched = true;
            if (OPT_no_trim)
            {
                fh_out[primer_group_name]->write_pair(seq1, seq2, PE_STRAND_OBJ);
            }
            else
            {
                htio2::FastqSeq trimmed1;
                htio2::FastqSeq trimmed2;
                seq1.subseq(trimmed1, trim_fwd_r1);
                seq2.subseq(trimmed2, trim_fwd_r2);
                fh_out[primer_group_name]->write_pair(trimmed1, trimmed2, PE_STRAND_OBJ);
            }
        }
        // revcom ok
        else if ((!n_primer_left || left_rev_ok)
                 &&
                 (!n_primer_right || right_rev_ok))
        {
            n_accept++;
            matched = true;
            if (OPT_no_trim)
            {
                fh_out[primer_group_name]->write_pair(seq1, seq2, PE_STRAND_OBJ);
            }
            else
            {
                htio2::FastqSeq trimmed1;
                htio2::FastqSeq trimmed2;
                seq1.subseq(trimmed1, trim_rev_r1);
                seq2.subseq(trimmed2, trim_rev_r2);
                fh_out[primer_group_name]->write_pair(trimmed1, trimmed2, PE_STRAND_OBJ);
            }
        }

        if (matched)
            break;
    }

    if (!matched)
        fh_fail->write_pair(seq1, seq2);

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

void trim_pe()
{
    htio2::FastqSeq seq1;
    htio2::FastqSeq seq2;

    // output handles
    map<string, htio2::PairedEndIO::Ptr> fh_out_map;
    for (map<string, PrimerGroup>::const_iterator i = primers.begin(); i != primers.end(); i++)
    {
        const string& name = i->first;
        File prefix_out = File(OPT_dir_out).getChildFile(String(name));

        if (OPT_pe_itlv)
        {
            string file_out;
            add_fastq_suffix_se(prefix_out.getFullPathName().toStdString(), OPT_compress, file_out);
            check_output_overwrite(OPT_files_in, file_out);

            htio2::PairedEndIO::Ptr fh(new htio2::PairedEndIO(file_out, htio2::WRITE, PE_STRAND_FILE, OPT_compress));
            fh_out_map.insert(make_pair(name, fh));
        }
        else
        {
            string file_out1;
            string file_out2;
            add_fastq_suffix_pe(prefix_out.getFullPathName().toStdString(), OPT_compress, file_out1, file_out2);
            check_output_overwrite(OPT_files_in, file_out1);
            check_output_overwrite(OPT_files_in, file_out2);

            htio2::PairedEndIO::Ptr fh(new htio2::PairedEndIO(file_out1, file_out2, htio2::WRITE, PE_STRAND_FILE, OPT_compress));
            fh_out_map.insert(make_pair(name, fh));
        }
    }

    // handle for negative reads
    htio2::PairedEndIO::Ptr fh_fail;
    {
        File prefix_fail = File(OPT_dir_out).getChildFile("PRIMER_FAIL");
        if (OPT_pe_itlv)
        {
            string file_fail;
            add_fastq_suffix_se(prefix_fail.getFullPathName().toStdString(), OPT_compress, file_fail);
            check_output_overwrite(OPT_files_in, file_fail);
            fh_fail = new htio2::PairedEndIO(file_fail, htio2::WRITE, PE_STRAND_FILE, OPT_compress);
        }
        else
        {
            string file_fail1;
            string file_fail2;
            add_fastq_suffix_pe(prefix_fail.getFullPathName().toStdString(), OPT_compress, file_fail1, file_fail2);
            check_output_overwrite(OPT_files_in, file_fail1);
            check_output_overwrite(OPT_files_in, file_fail2);

            fh_fail = new htio2::PairedEndIO(file_fail1, file_fail2, htio2::WRITE, PE_STRAND_FILE, OPT_compress);
        }
    }

    // input handle
    htqc::MultiSeqFilePE* fh_in;
    if (OPT_pe_itlv)
    {
        fh_in = new htqc::MultiSeqFilePE(OPT_files_in, PE_STRAND_FILE, OPT_compress);
    }
    else
    {
        vector<string> files_in_a;
        vector<string> files_in_b;
        separate_paired_files(OPT_files_in, files_in_a, files_in_b);
        fh_in = new htqc::MultiSeqFilePE(files_in_a, files_in_b, PE_STRAND_FILE, OPT_compress);
    }

    // traverse input file pairs
    while (fh_in->next_pair(seq1, seq2, PE_STRAND_OBJ))
        do_one_pe(seq1, seq2, fh_out_map, fh_fail.get());

    delete fh_in;

    if (!OPT_quiet)
        cout << "done: " << n_seq << " pairs, " << n_accept << " trimmed" << endl;
}

void execute()
{
    read_primers();

    File dir_out_obj(OPT_dir_out);
    if (!dir_out_obj.isDirectory())
    {
        Result re = dir_out_obj.createDirectory();
        if (re.failed())
        {
            fprintf(stderr, "ERROR: failed to create output directory \"%s\": %s\n",
                    OPT_dir_out.c_str(), re.getErrorMessage().toRawUTF8());
        }
    }

    if (OPT_se)
        trim_se();
    else if (OPT_pe)
        trim_pe();
    else
        throw runtime_error("neither single-end nor paired-end");
}
