//     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 <vector>
#include <map>
#include <string>
#include <iostream>
#include <fstream>
#include <sstream>
#include <stdexcept>
#include <algorithm>
#include <cmath>
#include <queue>

#include <cstdio>

#include "htio2/Cast.h"
#include "htio2/FastqIO.h"
#include "htio2/QualityUtil.h"
#include "htio2/HeaderUtil.h"

#define USE_se_pe
#define USE_encode
#define USE_header_format
#define USE_mask

#include "htio2/RingBuffer.h"
#include "htqc/Options.h"
#include "htqc/BaseQualCounter.h"
#include "htqc/FastqStat.h"
#include "htqc/App.h"
#include "htqc/MultiSeqFile.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;
using namespace htio2::juce;

//
// thread devices
//
struct StatWorker;
vector<StatWorker*> WORKERS;

queue<ssize_t> IDX_FREE;
queue<ssize_t> IDX_FILLED;

// list of reads, one for each thread
vector<htio2::FastqSeq> READ_BUFFER_1;
vector<htio2::FastqSeq> READ_BUFFER_2;

size_t N_PROCESSED = 0;

//
// command-line arguments
//
int OPT_num_threads = 2;
htio2::Option OPT_num_threads_ENTRY("threads", 't', OPTION_GROUP_MISC,
                                    &OPT_num_threads, htio2::Option::FLAG_NONE,
                                    "Number of worker threads.", "INT");

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");

//
// thread body
//

struct StatWorker: public Thread
{

    StatWorker(int thread_id)
        : Thread(String(thread_id))
        , todo(OPT_num_threads * 2)
        , done(OPT_num_threads * 2)
        , stat_a(OPT_mask, OPT_encode, OPT_header_format, OPT_header_sra)
        , stat_b(OPT_mask, OPT_encode, OPT_header_format, OPT_header_sra)
    {
    }

    virtual void run()
    {
        while (1)
        {
            ssize_t i = numeric_limits<ssize_t>::min();

            // get indices of the todo read
            todo.pop_sync(i);

            // exit when get the signal
            if (i == -1) return;


            // process fetched ones
            if (OPT_se)
            {
                stat_a.record_seq(READ_BUFFER_1[i]);
            }
            else
            {
                stat_a.record_seq(READ_BUFFER_1[i]);
                stat_b.record_seq(READ_BUFFER_2[i]);
            }

            // push back processed read slot
            done.push_sync(i);
        }
    }

    htio2::RingBuffer<ssize_t> todo;
    htio2::RingBuffer<ssize_t> done;
    FastqStat stat_a;
    FastqStat stat_b;
};

//
// functions
//

void write_cycle_quality_table(const File& file, const FastqStat& stat)
{
    ofstream OUT(file.getFullPathName().toRawUTF8());

    // write header
    OUT << "#cycle";
    for (int q = stat.qual_from; q <= stat.qual_to; q++)
        OUT << "\t" << q;
    OUT << endl;

    // write qualities for each cycle
    for (int i = 0; i < stat.observed_max_len; i++)
    {
        const BaseQualCounter& quals = stat.base_qual_counter[i];
        OUT << i + 1;
        for (int q = stat.qual_from; q <= stat.qual_to; q++)
            OUT << "\t" << quals.get_num(q);
        OUT << endl;
    }
}

void write_cycle_composition(const File& file, const FastqStat& stat, bool qual_mask)
{
    ofstream OUT(file.getFullPathName().toRawUTF8());

    // write header
    OUT << "cycle" << "\t" << "A" << "\t" << "T" << "\t" << "G" << "\t" << "C" << "\t" << "N";
    if (qual_mask) OUT << "\t" << "mask";
    OUT << endl;

    // write composition for each cycle
    for (int i = 0; i < stat.observed_max_len; i++)
    {
        const BaseCounter& counter = stat.base_compo_counter[i];
        OUT << i + 1 << "\t"
            << counter.a << "\t"
            << counter.t << "\t"
            << counter.g << "\t"
            << counter.c << "\t"
            << counter.n << "\t";
        if (qual_mask) OUT << counter.mask;
        OUT << endl;
    }
}

void write_se_seq_quality(const File& file, const FastqStat& stat)
{
    ofstream OUT(file.getFullPathName().toRawUTF8());

    // write header
    OUT << "#quality\tread 1 num\tread 1 accumulate" << endl;

    int q = stat.qual_from;
    int val_a = 0;
    int sum_a = 0;

    while (1)
    {
        if (q > stat.qual_to) break;

        val_a = stat.read_qual_counter.get_num(q);
        sum_a += val_a;
        OUT << q << "\t" << val_a << "\t" << sum_a << endl;

        q++;
    }

}

void write_pe_seq_quality(const File& file, const FastqStat& stat_a, const FastqStat& stat_b)
{
    ofstream OUT(file.getFullPathName().toRawUTF8());

    // write header
    OUT << "#quality\tread 1 num\tread 1 accumulate\tread 2 num\tread 2 accumulate" << endl;

    int q = stat_a.qual_from;
    int val_a = 0;
    int sum_a = 0;
    int val_b = 0;
    int sum_b = 0;

    while (1)
    {
        if (q > stat_a.qual_to) break;

        val_a = stat_a.read_qual_counter.get_num(q);
        val_b = stat_b.read_qual_counter.get_num(q);
        sum_a += val_a;
        sum_b += val_b;

        OUT << q << "\t" << val_a << "\t" << sum_a << "\t" << val_b << "\t" << sum_b << endl;

        q++;
    }
}

void write_se_seq_length(const File& file, const FastqStat& stat)
{
    ofstream OUT(file.getFullPathName().toRawUTF8());

    // write header
    OUT << "#length" << "\t" << "read 1 num" << "\t" << "read 1 accumulate" << endl;

    {
        int i = 1;
        int val_a = 0;
        int sum_a = 0;

        while (1)
        {
            if (i > stat.observed_max_len) break;

            val_a = stat.read_len_counter[i];
            sum_a += val_a;
            OUT << i << "\t" << val_a << "\t" << sum_a << endl;

            i++;
        }
    }
}

void write_pe_seq_length(const File& file, const FastqStat& stat_a, const FastqStat& stat_b)
{
    ofstream OUT(file.getFullPathName().toRawUTF8());

    // write header
    if (stat_a.observed_max_len != stat_b.observed_max_len)
    {
        cerr << "ERROR: store has inequal length: " << stat_a.observed_max_len << " and " << stat_b.observed_max_len << endl;
        exit(EXIT_FAILURE);
    }

    size_t length = stat_a.observed_max_len;

    OUT << "#length" << "\t" << "read 1 num" << "\t" << "read 1 accumulate" << "\t" << "read 2 num" << "\t" << "read 2 accumulate" << endl;

    {
        int i = 1;
        int val_a = 0;
        int sum_a = 0;
        int val_b = 0;
        int sum_b = 0;

        while (1)
        {
            if (i > length) break;

            val_a = stat_a.read_len_counter[i];
            val_b = stat_b.read_len_counter[i];
            sum_a += val_a;
            sum_b += val_b;

            OUT << i << "\t" << val_a << "\t" << sum_a << "\t" << val_b << "\t" << sum_b << endl;

            i++;
        }
    }
}

void write_lane_tile_quality(const File& file, const FastqStat& stat)
{
    ofstream OUT(file.getFullPathName().toRawUTF8());

    // write header
    OUT << "#lane\ttile\tqual<10\t10<qual<20\t20<qual<30\tqual>30" << endl;

    // write content
    for (
         map< int, map<int, TileQualCounter> >::const_iterator it_lane = stat.tile_read_qual_counter.begin();
         it_lane != stat.tile_read_qual_counter.end();
         it_lane++
         )
    {
        const map<int, TileQualCounter>* lane_quals = &(it_lane->second);

        for (
             map<int, TileQualCounter>::const_iterator it_tile = lane_quals->begin();
             it_tile != lane_quals->end();
             it_tile++
             )
        {
            const TileQualCounter* tile_quals = &(it_tile->second);
            OUT << it_lane->first << "\t"
                << it_tile->first << "\t"
                << tile_quals->num_reads_10 << "\t"
                << tile_quals->num_reads_10_20 << "\t"
                << tile_quals->num_reads_20_30 << "\t"
                << tile_quals->num_reads_30 << endl;
        }
    }
}

void write_qq(const File& file, const FastqStat& stat_a, const FastqStat& stat_b)
{
    ofstream OUT(file.getFullPathName().toRawUTF8());

    // calculate pearson correlation
    const size_t n = stat_a.read_quals.size();

    double sum_xy(0);
    double sum_x(0);
    double sum_y(0);
    double sum_xx(0);
    double sum_yy(0);

    for (size_t i = 0; i < n; i++)
    {
        float x = stat_a.read_quals[i];
        float y = stat_b.read_quals[i];
        sum_xy += x*y;
        sum_x += x;
        sum_y += y;
        sum_xx += x*x;
        sum_yy += y*y;
    }

    double pearson =
            (
                n * sum_xy - sum_x * sum_y
                ) / (
                sqrt(n * sum_xx - sum_x * sum_x) * sqrt(n * sum_yy - sum_y * sum_y)
                );

    vector<float> quals_a_sorted = stat_a.read_quals;
    vector<float> quals_b_sorted = stat_b.read_quals;

    sort(quals_a_sorted.begin(), quals_a_sorted.end());
    sort(quals_b_sorted.begin(), quals_b_sorted.end());
    int num_seq = stat_a.read_quals.size();
    int step = (num_seq < 200 ? 1 : num_seq / 200);

    OUT << "# pearson correlation: " << pearson << endl;

    // write data
    for (int i = 0; i < num_seq; i += step)
    {
        OUT << quals_a_sorted[i] << "\t" << quals_b_sorted[i] << endl;
    }
}

//
// main
//

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_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_header_format_ENTRY);
    opt_psr.add_option(OPT_header_sra_ENTRY);

    opt_psr.add_option(OPT_num_threads_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] << " - summarize sequencing reads quality"
             << 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_header_format();

    if ((OPT_encode == htio2::ENCODE_SOLEXA) && OPT_mask)
    {
        cerr << "ERROR: solexa mode cannot be used together with mask" << endl;
        exit(EXIT_FAILURE);
    }
}

void show_options()
{
    show_se_pe();
    show_files_in();
    cout << "# output directory: " << OPT_dir_out << endl;
    show_encode();
    show_header_format();
    show_mask();
}

void send_all_filled_slots()
{
    while (IDX_FILLED.size())
    {
        // get a filled slot
        ssize_t slot_i = IDX_FILLED.front();
        IDX_FILLED.pop();

        // send this slot to a worker
        bool slot_sent = false;
        while (!slot_sent)
        {
            for (int i_worker = 0; i_worker < OPT_num_threads; i_worker++)
            {
                slot_sent = WORKERS[i_worker]->todo.push(slot_i);
                if (slot_sent) break;
            }
        }

        // log
        N_PROCESSED++;
        if (!OPT_quiet && (N_PROCESSED % LOG_BLOCK == 0))
            cout << "\t" << N_PROCESSED << " reads done" << endl;
    }
}

void fetch_empty_slots()
{
    for (int i_worker = 0; i_worker < OPT_num_threads; i_worker++)
    {
        ssize_t slot_i;
        if (WORKERS[i_worker]->done.pop(slot_i))
        {
            IDX_FREE.push(slot_i);
        }
        if (WORKERS[i_worker]->done.pop(slot_i))
        {
            IDX_FREE.push(slot_i);
        }
    }
}

void do_single()
{
    htio2::FastqSeq curr;

    htqc::MultiSeqFileSE fh_in(OPT_files_in, OPT_compress);

    bool fh_re = false;

    for (;;)
    {
        while (IDX_FREE.size())
        {
            ssize_t slot_i = IDX_FREE.front();
            IDX_FREE.pop();

            fh_re = fh_in.next_seq(READ_BUFFER_1[slot_i]);
            if (!fh_re)
                break;

            IDX_FILLED.push(slot_i);
        }

        // send all filled slots to worker
        send_all_filled_slots();

        // return if finished
        if (!fh_re)
            break;

        // try to fetch some empty slots
        fetch_empty_slots();
    }

    if (!OPT_quiet)
        cout << "\t" << N_PROCESSED << " done" << endl;
}


void do_paired()
{
    htqc::MultiSeqFilePE::Ptr 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);
    }

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

    bool fh_re = false;

    for (;;)
    {
        // read from file
        while (IDX_FREE.size())
        {
            ssize_t slot_i = IDX_FREE.front();
            IDX_FREE.pop();

            fh_re = fh_in->next_pair(READ_BUFFER_1[slot_i], READ_BUFFER_2[slot_i], PE_STRAND_OBJ);
            if (!fh_re)
                break;

            IDX_FILLED.push(slot_i);
        }

        // send all filled slots to worker
        send_all_filled_slots();

        // return if finished
        if (!fh_re)
            break;

        // try to read empty slots
        fetch_empty_slots();
    }

    if (!OPT_quiet)
        cout << "\t" << N_PROCESSED << " pairs done" << endl;
}

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

    //
    // create thread communication utilities
    //
    READ_BUFFER_1.resize(OPT_num_threads*2);
    READ_BUFFER_2.resize(OPT_num_threads*2);

    // mark all slots as ready to use
    for (ssize_t i = 0; i < OPT_num_threads*2; i++)
    {
        IDX_FREE.push(i);
    }

    // create threads
    for (int i = 0; i < OPT_num_threads; i++)
    {
        WORKERS.push_back(new StatWorker(i));
        WORKERS.back()->startThread();
    }

    //
    // traverse input
    //
    if (OPT_se)
        do_single();
    else if (OPT_pe)
        do_paired();
    else
        throw runtime_error("neither single-end nor paired-end");

    // send stop signal
    for (int i = 0; i < OPT_num_threads; i++)
    {
        WORKERS[i]->todo.push_sync(-1);
    }

    // join threads
    // reduce statistic results from all threads into one
    FastqStat stat_a(OPT_mask, OPT_encode, OPT_header_format, OPT_header_sra);
    FastqStat stat_b(OPT_mask, OPT_encode, OPT_header_format, OPT_header_sra);

    if (!OPT_quiet) cout << "join threads" << endl;
    for (int i = 0; i < OPT_num_threads; i++)
    {
        StatWorker* worker = WORKERS[i];
        if (!OPT_quiet) cout << "  thread " << i << endl;
        worker->waitForThreadToExit(-1);

        if (worker->stat_a.observed_max_len < stat_a.observed_max_len)
            worker->stat_a.try_expand_store(stat_a.observed_max_len);
        stat_a += worker->stat_a;

        if (!OPT_se)
        {
            if (worker->stat_b.observed_max_len < stat_b.observed_max_len)
                worker->stat_b.try_expand_store(stat_b.observed_max_len);
            stat_b += worker->stat_b;
        }

        delete WORKERS[i];
    }

    if (!OPT_se)
    {
        if (stat_a.observed_max_len > stat_b.observed_max_len)
            stat_b.try_expand_store(stat_a.observed_max_len);
        else
            stat_a.try_expand_store(stat_b.observed_max_len);
    }

    //
    // write result
    //
    if (!OPT_quiet) cout << "write result" << endl;

    // reads info
    File file_info = dir_out.getChildFile("info.tab");
    ofstream INFO(file_info.getFullPathName().toRawUTF8());
    INFO << "amount" << "\t" << N_PROCESSED << (OPT_se ? "" : "x2") << endl
         << "encode" << "\t" << htio2::to_string(OPT_encode) << endl
         << "paired" << "\t" << bool_to_string(OPT_pe) << endl
         << "length" << "\t" << stat_a.observed_max_len << endl
         << "mask" << "\t" << bool_to_string(OPT_mask) << endl
         << "quality range" << "\t" << stat_a.qual_from << '-' << stat_a.qual_to << endl;

    // cycle quality
    File file_cycle_qual_1 = dir_out.getChildFile("cycle_quality_1.tab");
    if (!OPT_quiet)
        cout << "  " << file_cycle_qual_1.getFullPathName().toStdString() << endl;

    write_cycle_quality_table(file_cycle_qual_1, stat_a);

    if (!OPT_se)
    {
        File file_cyc_qual_2 = dir_out.getChildFile("cycle_quality_2.tab");
        if (!OPT_quiet)
            cout << "  " << file_cyc_qual_2.getFullPathName().toStdString() << endl;

        write_cycle_quality_table(file_cyc_qual_2, stat_b);
    }

    // read quality
    File file_read_qual = dir_out.getChildFile("reads_quality.tab");
    if (!OPT_quiet) cout << "  " << file_read_qual.getFullPathName().toStdString() << endl;

    if (!OPT_se)
    {
        write_pe_seq_quality(file_read_qual, stat_a, stat_b);
    }
    else
    {
        write_se_seq_quality(file_read_qual, stat_a);
    }

    // lane and tile quality
    File file_lane_tile_qual_1 = dir_out.getChildFile("lane_tile_quality_1.tab");
    if (!OPT_quiet)
        cout << "  " << file_lane_tile_qual_1.getFullPathName().toStdString() << endl;
    write_lane_tile_quality(file_lane_tile_qual_1, stat_a);

    if (!OPT_se)
    {
        File file_lane_tile_qual_2 = dir_out.getChildFile("lane_tile_quality_2.tab");
        if (!OPT_quiet)
            cout << "  " << file_lane_tile_qual_2.getFullPathName().toStdString() << endl;
        write_lane_tile_quality(file_lane_tile_qual_2, stat_b);
    }

    // cycle composition
    File file_cycle_comp_1 = dir_out.getChildFile("cycle_composition_1.tab");
    if (!OPT_quiet)
        cout << "  " << file_cycle_comp_1.getFullPathName().toStdString() << endl;
    write_cycle_composition(file_cycle_comp_1, stat_a, OPT_mask);

    if (!OPT_se)
    {
        File file_cycle_comp_2 = dir_out.getChildFile("cycle_composition_2.tab");
        if (!OPT_quiet) cout << "  " << file_cycle_comp_2.getFullPathName().toStdString() << endl;
        write_cycle_composition(file_cycle_comp_2, stat_b, OPT_mask);
    }

    // read length
    File file_read_length = dir_out.getChildFile("reads_length.tab");
    if (!OPT_quiet) cout << "  " << file_read_length.getFullPathName().toStdString() << endl;

    if (!OPT_se)
    {
        write_pe_seq_length(file_read_length, stat_a, stat_b);
    }
    else
    {
        write_se_seq_length(file_read_length, stat_a);
    }

    // QQ plot
    if (!OPT_se)
    {
        File file_qual_qq = dir_out.getChildFile("quality_QQ.tab");
        if (!OPT_quiet) cout << "  " << file_qual_qq.getFullPathName().toStdString() << endl;
        write_qq(file_qual_qq, stat_a, stat_b);
    }

    if (!OPT_quiet)
        printf("all table written\n");
}

