//     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/PairedEndIO.h"

#define USE_prefix_out
#define USE_se_pe
#define USE_interleave
#include "htqc/Options.h"
#include "htqc/App.h"
#include "htqc/MultiSeqFile.h"
#include "htio2/MT19937.h"

#include <set>
#include <iostream>
#include <fstream>

// We don't care strand in this program, so simply let FILE and OBJ to be 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;
using namespace htio2::juce;

#define OPTION_GROUP_SAMPLING "Sampling Parameters"

double OPT_rate = 0;
htio2::Option OPT_rate_ENTRY("rate", 'r', OPTION_GROUP_SAMPLING,
                             &OPT_rate, htio2::Option::FLAG_NONE,
                             "1/rate of total reads will be picked.", "NUMBER");

size_t OPT_num = 0;
htio2::Option OPT_num_ENTRY("num", 'n', OPTION_GROUP_SAMPLING,
                            &OPT_num, htio2::Option::FLAG_NONE,
                            "Number of reads to be picked.", "INT");

int64_t OPT_seed = 0;
htio2::Option OPT_seed_ENTRY("seed", '\0', OPTION_GROUP_SAMPLING,
                             &OPT_seed, htio2::Option::FLAG_NONE,
                             "Seed for random number generator. 0 to use system-random seed.", "INT");

string OPT_file_out_id;
htio2::Option OPT_file_out_id_ENTRY("out-id", '\0', OPTION_GROUP_IO,
                                    &OPT_file_out_id, htio2::Option::FLAG_NONE,
                                    "Write the indices of picked reads to a file. The indices are 0-based.", "FILE");



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_file_out_id_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_rate_ENTRY);
    opt_psr.add_option(OPT_num_ENTRY);
    opt_psr.add_option(OPT_seed_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] << " - randomly pick some 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_prefix_out();

    if (OPT_rate == 0)
    {
        if (OPT_num == 0)
        {
            cerr << "ERROR: neither sampling rate nor sampling amount is specified." << endl
                 << "ERROR: you must specify how much to be sampled." << endl;
            exit(EXIT_FAILURE);
        }
    }
    else
    {
        if (OPT_num)
        {
            cerr << "ERROR: both sampling rate and sampling amount is specified." << endl
                 << "ERROR: you must specify only one of them." << endl;
            exit(EXIT_FAILURE);
        }
        if (OPT_rate <= 1)
        {
            cerr << "ERROR: invalid sampling rate: " << OPT_rate << endl
                 << "ERROR: sampling rate must larger than 1" << endl;
            exit(EXIT_FAILURE);
        }
    }
}

void show_options()
{
    show_se_pe();
    show_files_in();
    show_prefix_out();
    if (OPT_file_out_id.size())
        cout << "# output file for used indices: " << OPT_file_out_id << endl;

    if (OPT_rate > 0)
        cout << "# fraction to be picked: 1/" << OPT_rate << endl;
    if (OPT_num > 0)
        cout << "# amount to be picked: " << OPT_num << endl;
    if (OPT_seed)
        cout << "# user random seed: " << OPT_seed << endl;
}

int64_t n_accept = 0;
int64_t n_input = 0;

int64_t count_se()
{
    htqc::MultiSeqFileSE fh_in(OPT_files_in, OPT_compress);
    htio2::FastqSeq seq;
    int64_t n = 0;

    while (fh_in.next_seq(seq))
        n++;

    if (n < 0)
    {
        fprintf(stderr, "ERROR: sequence amount overflow!\n");
        exit(EXIT_FAILURE);
    }
    return n;
}

int64_t count_pe()
{
    htqc::MultiSeqFilePE* fh = nullptr;
    if (OPT_pe_itlv)
    {
        fh = 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 = new htqc::MultiSeqFilePE(files_in_a,
                                      files_in_b,
                                      PE_STRAND_FILE,
                                      OPT_compress);
    }

    htio2::FastqSeq seq1;
    htio2::FastqSeq seq2;
    int64_t n = 0;

    while (fh->next_pair(seq1, seq2, PE_STRAND_OBJ))
        n++;

    if (n < 0)
    {
        fprintf(stderr, "ERROR: sequence amount overflow!\n");
        exit(EXIT_FAILURE);
    }

    delete fh;
    return n;
}



void shuffle_vector(vector<int64_t>& data, htio2::MT19937& prng)
{
    for (size_t i1 = 0; i1 < data.size(); i1++)
    {
        uint64_t i2 = prng.next_uint64_in_range(data.size());
        printf("  switch %lu %lu\n", i1, i2);
        int64_t tmp = data[i1];
        data[i1] = data[i2];
        data[i2] = tmp;
    }
}

void rand_pick(int64_t num, set<int64_t>& wanted, htio2::MT19937& prng)
{
    vector<int64_t> indices(num);
    for (int64_t i = 0; i < num; i++)
        indices[i] = i;

    // shuffle
    shuffle_vector(indices, prng);

    // do pick
    size_t n_pick = 0;
    if (OPT_num)
        n_pick = OPT_num;
    else
        n_pick = indices.size() / OPT_rate;

    for (int i = 0; i < n_pick; i++)
    {
        if (!wanted.insert(indices[i]).second)
            abort();
    }
}

void filter_se(const set<int64_t>& wanted)
{
    // 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 OUT(htio2::SeqIO::New(file_out,
                                          htio2::WRITE,
                                          OPT_compress));

    // input handle
    htqc::MultiSeqFileSE fh_in(OPT_files_in, OPT_compress);

    // traverse input sequences
    htio2::FastqSeq curr_seq;
    int64_t id = 0;

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

        if (wanted.count(id++))
        {
            OUT->write_seq(curr_seq);
            n_accept++;
        }

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



    if (!OPT_quiet)
        cout << "  " << n_input << " reads, " << n_accept << " accept" << endl;
}

void filter_pe(const set<int64_t>& wanted)
{
    htio2::PairedEndIO* fh_out = NULL;
    htqc::MultiSeqFilePE* fh_in = NULL;

    if (OPT_pe_itlv)
    {
        string file_out;
        add_fastq_suffix_se(OPT_prefix_out, OPT_compress, file_out);
        check_output_overwrite(OPT_files_in, file_out);

        fh_out = new htio2::PairedEndIO(file_out,
                                        htio2::WRITE,
                                        PE_STRAND_FILE,
                                        OPT_compress);

        fh_in = new htqc::MultiSeqFilePE(OPT_files_in,
                                         PE_STRAND_FILE,
                                         OPT_compress);
    }
    else
    {
        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);

        fh_out = new htio2::PairedEndIO(file_out_a,
                                        file_out_b,
                                        htio2::WRITE,
                                        PE_STRAND_FILE,
                                        OPT_compress);

        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 sequences
    htio2::FastqSeq curr_seq_a;
    htio2::FastqSeq curr_seq_b;
    int64_t id = 0;

    while (fh_in->next_pair(curr_seq_a, curr_seq_b, PE_STRAND_OBJ))
    {
        n_input++;

        if (wanted.count(id++))
        {
            fh_out->write_pair(curr_seq_a, curr_seq_b, PE_STRAND_OBJ);
            n_accept++;
        }

        if (!OPT_quiet && (n_input % LOG_BLOCK == 0))
            cout << "\t" << n_input << " pairs, " << n_accept << " picked" << endl;
    }

    delete fh_in;
    delete fh_out;

    if (!OPT_quiet)
        cout << "  " << n_input << " pairs, " << n_accept << " picked" << endl;
}

void execute()
{
    // init PRNG
    htio2::MT19937 prng;
    if (OPT_seed != 0)
        prng.set_seed(OPT_seed);

    // first pass
    // get sequence num
    if (!OPT_quiet)
        printf("count input sequences\n");

    int64_t num = 0;
    if (OPT_se)
        num = count_se();
    else if (OPT_pe)
        num = count_pe();
    else
        abort();

    if (!OPT_quiet)
        printf("  %ld total\n", num);

    // generate a table of wanted indices
    if (!OPT_quiet)
        printf("randomly generate wanted index table\n");
    set<int64_t> wanted;
    rand_pick(num, wanted, prng);

    if (!OPT_quiet)
        printf("  %lu indices\n", wanted.size());

    // second pass
    // pick wanted sequences
    if (OPT_se)
        filter_se(wanted);
    else if (OPT_pe)
        filter_pe(wanted);
    else
        abort();

    // write a list of used indices
    if (OPT_file_out_id.length())
    {
        if (!OPT_quiet)
            printf("write used indices to file \"%s\"\n", OPT_file_out_id.c_str());

        ofstream fh_out_id(OPT_file_out_id.c_str());
        for (int64_t id : wanted)
        {
            fh_out_id << id << endl;
        }
    }
}
